import tf
import gc
import pyb
import time
import sensor
import machine
import micropython

from micropython import const
from uart import uart_init, send_data, recv_data

# -------------------------------程序初始化部分----------------------------

# 摄像头初始化
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QQVGA)
sensor.set_windowing((120, 120))
# sensor.set_framebuffers(sensor.DOUBLE_BUFFER)
sensor.set_framebuffers(sensor.TRIPLE_BUFFER)
sensor.set_contrast(3)
sensor.set_saturation(3)
sensor.skip_frames(n=100)

img = sensor.snapshot()
clock = time.clock()

# 常量定义
WIDTH = const(320)
HEIGHT = const(240)
COLORED_THRESHOLD = const((0, 255, -30, 30, -30, 30))
STDEV_THRESHOLD = const(8)
CORR_SIZE = const(56)
COLOR_NAME = const(("Red", "Green", "Yellow", "Blue"))

# 串口初始化
huart = pyb.UART(3, 9600, timeout_char=200)
led = pyb.LED(1)
usb = pyb.USB_VCP()

# 初始化变量
detect = 1
rota_corr = 0

# -------------------------------程序主体部分----------------------------

@micropython.native
def led_ready(_led):
    """
    LED指示函数，用于在程序开始时点亮LED
    参数:
        _led: LED对象
    """
    _led.on()
    time.sleep_ms(300)
    _led.off()
    time.sleep_ms(300)
    _led.on()
    time.sleep_ms(300)
    _led.off()
    time.sleep_ms(300)


@micropython.native
def num_color(_stat):
    """
    根据图像统计结果判断颜色
    参数:
        _stat: 图像统计结果
    返回:
        颜色索引
    """
    _l = [
        _stat.a_max(),
        abs(_stat.a_min()),
        _stat.b_max(),
        abs(_stat.b_min()),
    ]

    # color_name = ["Red", "Green", "Yellow", "Blue"]
    return _l.index(max(_l))


@micropython.native
def num_calc_center(_img):
    """
    计算数字在图像中的中心点
    参数:
        _img: 图像对象
    返回:
        数字中心点的x坐标, y坐标和像素点数量
    """
    _x_avg = 0
    _y_avg = 0
    _count = 0

    for y, _ in enumerate(_img):
        for x, p in enumerate(_):
            if p > 200:
                _x_avg += x
                _y_avg += y
                _count += 1

    return _x_avg // _count, _y_avg // _count, _count


@micropython.native
def num_classify(_img):
    """
    使用训练好的模型对数字进行分类
    参数:
        _img: 图像对象
    返回:
        分类结果
    """
    _output = tf.classify(
        "trained.tflite",
        _img,
        min_scale=1.0,
        scale_mul=0.5,
        x_overlap=0.95,
        y_overlap=0.95,
    )[0].output()

    return _output.index(max(_output))


def loop():
    """
    主循环函数，负责图像处理和数字识别
    """
    global img, detect, rota_corr

    clock.tick()

    gc.collect()

    while not sensor.get_frame_available():
        pass

    len = huart.any()
    if len > 0:
        detect, rota_corr = recv_data(huart, len)

    img = sensor.snapshot()
    img_gray = img.copy().gaussian(1).binary([COLORED_THRESHOLD])

    # 判断并识别数字
    stdev = img_gray.get_statistics().stdev()
    if stdev > STDEV_THRESHOLD and detect:
        # 提取数字
        img.b_xor(img, mask=img_gray)

        # 计算图像中点
        img_gray.invert()
        x_avg, y_avg, count = num_calc_center(img_gray)

        # 矫正图像
        _corr = CORR_SIZE + (count // 43 - 28)
        corner = [
            (x_avg - _corr, y_avg - _corr),
            (x_avg + _corr, y_avg - _corr),
            (x_avg + _corr, y_avg + _corr),
            (x_avg - _corr, y_avg + _corr),
        ]

        img_gray.rotation_corr(z_rotation=rota_corr, corners=corner)
        img.rotation_corr(z_rotation=rota_corr, corners=corner)
        del _corr

        # 识别
        stat = img.get_statistics()
        color = num_color(stat)
        result = num_classify(img_gray)

        send_data(huart, result, color)
        if usb.isconnected():
            print(f"FPS:{clock.fps():.2f} NUM:{result} COLOR:{COLOR_NAME[color]}")
    elif usb.isconnected():
        print(f"FPS:{clock.fps():.2f}")

# 程序入口
led_ready(led)
del led

while True:
    try:
        loop()
    except Exception as e:
        print(e)
        machine.soft_reset()